from PIL import Image
import math


def canny_edge_detection(image_path, output_path):
    """
    使用Canny算子对输入图像进行边缘检测。

    参数:
    image_path (str): 输入图像的文件路径。
    output_path (str): 输出边缘检测后图像的文件路径。
    """
    # 打开原始图像并转换为灰度图（以灰度图为例）
    image = Image.open(image_path).convert('L')
    # 获取图像的宽度和高度
    width, height = image.size
    # 创建新的空白图像，尺寸与原图像相同，模式也相同，用于存储边缘检测结果
    new_image = Image.new('L', (width, height))

    # 1. 先用高斯滤波对图像进行消噪
    def gaussian_filter(image):
        """
        对图像进行简单的高斯滤波。
        """
        kernel_size = 3
        radius = kernel_size // 2
        gaussian_kernel = [[1 / 16, 2 / 16, 1 / 16],
                           [2 / 16, 4 / 16, 2 / 16],
                           [1 / 16, 2 / 16, 1 / 16]]
        filtered_image = Image.new('L', (width, height))
        for y in range(radius, height - radius):
            for x in range(radius, width - radius):
                sum_value = 0
                for ky in range(-radius, radius + 1):
                    for kx in range(-radius, radius + 1):
                        pixel = image.getpixel((x + kx, y + ky))
                        sum_value += pixel * gaussian_kernel[ky + radius][kx + radius]
                filtered_image.putpixel((x, y), int(sum_value))
        return filtered_image

    filtered_image = gaussian_filter(image)

    # 2. 计算梯度幅值和方向
    gradient_x = [[0] * width for _ in range(height)]
    gradient_y = [[0] * width for _ in range(height)]
    gradient_magnitude = [[0] * width for _ in range(height)]
    gradient_direction = [[0] * width for _ in range(height)]
    sobel_x = [[-1, 0, 1],
               [-2, 0, 2],
               [-1, 0, 1]]
    sobel_y = [[-1, -2, -1],
               [0, 0, 0],
               [1, 2, 1]]
    for y in range(1, height - 1):
        for x in range(1, width - 1):
            # 计算水平方向梯度
            sum_x = 0
            for ky in range(3):
                for kx in range(3):
                    pixel = filtered_image.getpixel((x + kx - 1, y + ky - 1))
                    sum_x += pixel * sobel_x[ky][kx]
            gradient_x[y][x] = sum_x
            # 计算垂直方向梯度
            sum_y = 0
            for ky in range(3):
                for kx in range(3):
                    pixel = filtered_image.getpixel((x + kx - 1, y + ky - 1))
                    sum_y += pixel * sobel_y[ky][kx]
            gradient_y[y][x] = sum_y
            # 计算梯度幅值（采用平方根法）
            gradient_magnitude[y][x] = math.sqrt(sum_x ** 2 + sum_y ** 2)
            # 计算梯度方向（弧度制，转换到0-2*pi范围）
            gradient_direction[y][x] = math.atan2(sum_y, sum_x)
            if gradient_direction[y][x] < 0:
                gradient_direction[y][x] += 2 * math.pi

    # 3. 非极大值抑制
    def non_maximum_suppression(gradient_magnitude, gradient_direction):
        """
        对梯度幅值进行非极大值抑制。
        """
        suppressed_image = [[0] * width for _ in range(height)]
        for y in range(1, height - 1):
            for x in range(1, width - 1):
                angle = gradient_direction[y][x]
                if (0 <= angle < math.pi / 8) or (7 * math.pi / 8 <= angle <= 2 * math.pi):
                    neighbor_1 = gradient_magnitude[y][x - 1]
                    neighbor_2 = gradient_magnitude[y][x + 1]
                elif (math.pi / 8 <= angle < 3 * math.pi / 8):
                    neighbor_1 = gradient_magnitude[y - 1][x + 1]
                    neighbor_2 = gradient_magnitude[y + 1][x - 1]
                elif (3 * math.pi / 8 <= angle < 5 * math.pi / 8):
                    neighbor_1 = gradient_magnitude[y - 1][x]
                    neighbor_2 = gradient_magnitude[y + 1][x]
                else:
                    neighbor_1 = gradient_magnitude[y - 1][x - 1]
                    neighbor_2 = gradient_magnitude[y + 1][x + 1]
                if gradient_magnitude[y][x] >= neighbor_1 and gradient_magnitude[y][x] >= neighbor_2:
                    suppressed_image[y][x] = gradient_magnitude[y][x]
        return suppressed_image

    suppressed_image = non_maximum_suppression(gradient_magnitude, gradient_direction)

    # 4. 双阈值检测与边缘连接
    high_threshold = 100  # 高阈值，可根据图像情况调整
    low_threshold = 50  # 低阈值，可根据图像情况调整
    final_image = [[0] * width for _ in range(height)]
    for y in range(1, height - 1):
        for x in range(1, width - 1):
            if suppressed_image[y][x] >= high_threshold:
                final_image[y][x] = 255
            elif suppressed_image[y][x] >= low_threshold:
                # 检查8邻域内是否有强边缘（值为255的像素）
                has_strong_neighbor = False
                for ky in range(-1, 2):
                    for kx in range(-1, 2):
                        if suppressed_image[y + ky][x + kx] >= high_threshold:
                            has_strong_neighbor = True
                            break
                    if has_strong_neighbor:
                        break
                if has_strong_neighbor:
                    final_image[y][x] = 255
    for y in range(height):
        for x in range(width):
            new_image.putpixel((x, y), final_image[y][x])

    # 保存边缘检测后的图像
    new_image.save(output_path)
    return new_image

image_path = "D:/cangku/computer-image-project-design/test3.png"
output_path = "D:/cangku/computer-image-project-design/test6.png"
canny_edge_detection(image_path, output_path)